import { create } from "zustand";
import {
  type Node,
  type OnNodesChange,
  type Edge,
  type OnEdgesChange,
  type OnConnect,
  type Connection,
  applyNodeChanges,
  applyEdgeChanges,
  addEdge,
} from "reactflow";

interface SelectNodeProps {
  id?: string;
  select: boolean;
}

interface WorkflowState {
  nodes: Node[];
  edges: Edge[];
  selectedNode: Node | null;
  drawerVisible: boolean;
  addNode: (node: Node) => void;
  setNodes: (nodes: Node[]) => void;
  setEdges: (edges: Edge[]) => void;
  setSelectedNode: (node: SelectNodeProps) => void;
  updateNodeData: (node: Node) => void;
  setDrawerVisible: (drawerVisible: boolean) => void;
  onNodesChange: OnNodesChange;
  onEdgesChange: OnEdgesChange;
  onConnect: OnConnect;
  onReconnect: (oldEdge: Edge, newConnection: Connection) => void;
}

const useWorkflowStore = create<WorkflowState>()((set, get) => {
  return {
    nodes: [],
    edges: [],
    selectedNode: null,
    drawerVisible: false,
    setNodes: (nodes) => set({ nodes: nodes }),
    setEdges: (edges) => set({ edges: edges }),
    setSelectedNode: (node) => {
      set((state) => ({
        selectedNode: node
          ? state.nodes.find((n) => n.id === node.id) || null
          : null,
      }));
    },
    setDrawerVisible: (drawerVisible) => set({ drawerVisible: drawerVisible }),
    addNode: (node) => {
      set((state) => ({
        nodes: [...state.nodes, node],
      }));
    },
    updateNodeData: (node) => {
      set((state) => ({
        nodes: state.nodes.map((n) => {
          if (n.id === node.id) {
            return {
              ...n,
              data: { ...node.data },
            };
          }
          return n;
        }),
      }));
    },
    onNodesChange: (changes) => {
      for (const change of changes) {
        if (change.type === "select") {
          if (change.selected === true) {
            set({
              selectedNode: get().nodes.find((node) => node.id === change.id),
              drawerVisible: true,
            });
          } else {
            set({
              selectedNode: null,
              drawerVisible: false,
            });
          }
          if (change.selected === true && changes.length > 1) {
            break;
          }
        }
      }
      set({
        nodes: applyNodeChanges(changes, get().nodes),
      });
    },
    onEdgesChange: (changes) => {
      set({
        edges: applyEdgeChanges(changes, get().edges),
      });
    },
    onConnect: (params) => {
      const edge = {
        ...params,
        type: "customEdge",
        data: { hover: false },
      };
      set({
        edges: addEdge(edge, get().edges),
      });
    },
    onReconnect: (oldEdge: Edge, newConnection: Connection) => {
      const edge = {
        ...newConnection,
        type: "customEdge",
        data: {
          hover: false,
        },
      };
      set({
        edges: addEdge(
          edge,
          get().edges.filter((eg) => eg.id !== oldEdge.id)
        ),
      });
    },
  };
});

export default useWorkflowStore;
